package org.ruoyi.controller;

import cn.dev33.satoken.stp.StpUtil;
import jakarta.servlet.http.HttpServletResponse;
import jakarta.validation.Valid;
import jakarta.validation.constraints.NotEmpty;
import jakarta.validation.constraints.NotNull;
import lombok.Data;
import lombok.RequiredArgsConstructor;
import org.ruoyi.common.chat.config.ChatConfig;
import org.ruoyi.common.chat.domain.request.ChatRequest;
import org.ruoyi.common.chat.entity.chat.ChatCompletion;
import org.ruoyi.common.chat.entity.chat.Message;
import org.ruoyi.common.chat.openai.OpenAiStreamClient;
import org.ruoyi.common.core.domain.R;
import org.ruoyi.common.core.validate.AddGroup;
import org.ruoyi.common.excel.utils.ExcelUtil;
import org.ruoyi.common.log.annotation.Log;
import org.ruoyi.common.log.enums.BusinessType;
import org.ruoyi.common.mybatis.core.page.PageQuery;
import org.ruoyi.common.mybatis.core.page.TableDataInfo;
import org.ruoyi.common.satoken.utils.LoginHelper;
import org.ruoyi.common.web.core.BaseController;
import org.ruoyi.knowledge.domain.bo.KnowledgeAttachBo;
import org.ruoyi.knowledge.domain.bo.KnowledgeFragmentBo;
import org.ruoyi.knowledge.domain.bo.KnowledgeInfoBo;
import org.ruoyi.knowledge.domain.req.KnowledgeInfoUploadRequest;
import org.ruoyi.knowledge.domain.vo.KnowledgeAttachVo;
import org.ruoyi.knowledge.domain.vo.KnowledgeFragmentVo;
import org.ruoyi.knowledge.domain.vo.KnowledgeInfoVo;
import org.ruoyi.knowledge.service.EmbeddingService;
import org.ruoyi.knowledge.service.IKnowledgeAttachService;
import org.ruoyi.knowledge.service.IKnowledgeFragmentService;
import org.ruoyi.knowledge.service.IKnowledgeInfoService;
import org.ruoyi.knowledge.service.RagflowService;
import org.ruoyi.system.listener.SSEEventSourceListener;
import org.springframework.validation.annotation.Validated;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import org.ruoyi.knowledge.chain.vectorstore.VectorStore;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.HashMap;
import java.util.List;
import java.util.Map;


/**
 * 知识库
 *
 * @author Lion Li
 * @date 2024-10-21
 */
@Validated
@RequiredArgsConstructor
@RestController
@RequestMapping("/knowledge")
public class KnowledgeController extends BaseController {

    private final IKnowledgeInfoService knowledgeInfoService;

    private final VectorStore vectorStore;

    private final IKnowledgeAttachService attachService;

    private final IKnowledgeFragmentService fragmentService;

    private final EmbeddingService embeddingService;
    
    private final RagflowService ragflowService;

    private OpenAiStreamClient openAiStreamClient;

    private final ChatConfig chatConfig;

    /**
     * 知识库对话
     */
    @PostMapping("/send")
    public SseEmitter send(@RequestBody @Valid ChatRequest chatRequest) {
        openAiStreamClient = chatConfig.getOpenAiStreamClient();
        SseEmitter sseEmitter = new SseEmitter(0L);
        SSEEventSourceListener openAIEventSourceListener = new SSEEventSourceListener(sseEmitter);
        List<Message> messages = chatRequest.getMessages();
        String content = messages.get(messages.size() - 1).getContent().toString();
        List<String> nearestList;
        List<Double> queryVector = embeddingService.getQueryVector(content, chatRequest.getKid());
        nearestList = vectorStore.nearest(queryVector,chatRequest.getKid());
        for (String prompt : nearestList) {
            Message sysMessage = Message.builder().content(prompt).role(Message.Role.USER).build();
            messages.add(sysMessage);
        }
        Message userMessage = Message.builder().content(content + (nearestList.size() > 0 ? "\n\n注意：回答问题时，须严格根据我给你的系统上下文内容原文进行回答，请不要自己发挥,回答时保持原来文本的段落层级" : "") ).role(Message.Role.USER).build();
        messages.add(userMessage);
        ChatCompletion completion = ChatCompletion
            .builder()
            .messages(messages)
            .model(chatRequest.getModel())
            .temperature(chatRequest.getTemperature())
            .topP(chatRequest.getTop_p())
            .stream(true)
            .build();
        openAiStreamClient.streamChatCompletion(completion, openAIEventSourceListener);

        return sseEmitter;
    }

    /**
     * 根据用户信息查询本地知识库
     */
    @GetMapping("/list")
    public TableDataInfo<KnowledgeInfoVo> list(KnowledgeInfoBo bo, PageQuery pageQuery) {
        if(!StpUtil.isLogin()){
            throw new SecurityException("请先去登录!");
        }
        bo.setUid(LoginHelper.getUserId());
        return knowledgeInfoService.queryPageList(bo, pageQuery);
    }

    /**
     * 新增知识库
     */
    @Log(title = "知识库", businessType = BusinessType.INSERT)
    @PostMapping("/save")
    public R<Void> save(@Validated(AddGroup.class) @RequestBody KnowledgeInfoBo bo) {
        knowledgeInfoService.saveOne(bo);
        return R.ok();
    }

    /**
     * 删除知识库
     */
    @PostMapping("/remove/{id}")
    public R<String> remove(@PathVariable String id){
        knowledgeInfoService.removeKnowledge(id);
        return R.ok("删除知识库成功!");
    }

    /**
     * 修改知识库
     */
    @Log(title = "知识库", businessType = BusinessType.UPDATE)
    @PostMapping("/edit")
    public R<Void> edit( @RequestBody KnowledgeInfoBo bo) {
        return toAjax(knowledgeInfoService.updateByBo(bo));
    }

    /**
     * 导出知识库列表
     */
    @Log(title = "知识库", businessType = BusinessType.EXPORT)
    @PostMapping("/export")
    public void export(KnowledgeInfoBo bo, HttpServletResponse response) {
        List<KnowledgeInfoVo> list = knowledgeInfoService.queryList(bo);
        ExcelUtil.exportExcel(list, "知识库", KnowledgeInfoVo.class, response);
    }

    /**
     * 查询知识附件信息
     */
    @GetMapping("/detail/{kid}")
    public TableDataInfo<KnowledgeAttachVo> attach(KnowledgeAttachBo bo, PageQuery pageQuery,@PathVariable String kid){
        bo.setKid(kid);
        return attachService.queryPageList(bo, pageQuery);
    }

    /**
     * 上传知识库附件
     */
    @PostMapping(value = "/attach/upload")
    public R<String> upload(KnowledgeInfoUploadRequest request){
        knowledgeInfoService.upload(request);
        return R.ok("上传知识库附件成功!");
    }

    /**
     * 获取知识库附件详细信息
     *
     * @param id 主键
     */
    @GetMapping("attach/info/{id}")
    public R<KnowledgeAttachVo> getAttachInfo(@NotNull(message = "主键不能为空")
                                              @PathVariable Long id) {
        return R.ok(attachService.queryById(id));
    }

    /**
     * 删除知识库附件
     *
     */
    @PostMapping("attach/remove/{docId}")
    public R<Void> removeAttach(@NotEmpty(message = "主键不能为空") @PathVariable String docId) {
        attachService.removeKnowledgeAttach(docId);
        return R.ok();
    }


    /**
     * 查询知识片段（从本地数据库）
     */
    @GetMapping("/fragment/list/{docId}")
    public TableDataInfo<KnowledgeFragmentVo> fragmentList(KnowledgeFragmentBo bo, PageQuery pageQuery, @PathVariable String docId) {
        bo.setDocId(docId);
        return fragmentService.queryPageList(bo, pageQuery);
    }

    /**
     * 从RAGFlow获取文档的知识片段
     */
    @GetMapping("/fragment/ragflow/{docId}")
    public R<Map<String, Object>> getRagflowFragments(
            @PathVariable String docId,
            @RequestParam(defaultValue = "1") int pageNum,
            @RequestParam(defaultValue = "100") int pageSize) {
        try {
            // 获取附件信息
            KnowledgeAttachVo attach = attachService.queryById(Long.valueOf(docId));
            if (attach == null) {
                return R.fail("文档不存在");
            }
            
            // 获取知识库信息
            String kid = attach.getKid();
            List<KnowledgeInfoVo> knowledgeInfoList = knowledgeInfoService.queryList(
                new KnowledgeInfoBo() {{ setKid(kid); }}
            );
            if (knowledgeInfoList == null || knowledgeInfoList.isEmpty()) {
                return R.fail("知识库不存在");
            }
            
            String datasetId = knowledgeInfoList.get(0).getRagflowDatasetId();
            if (datasetId == null || datasetId.isEmpty()) {
                return R.fail("数据集ID为空");
            }
            
            // 通过文件名查找RAGFlow的document_id
            List<Map<String, Object>> documents = ragflowService.getDocuments(datasetId);
            
            System.out.println("========== 查找RAGFlow文档 ==========");
            System.out.println("本地文档名: " + attach.getDocName());
            System.out.println("数据集ID: " + datasetId);
            System.out.println("RAGFlow中的文档数量: " + (documents != null ? documents.size() : 0));
            
            String ragflowDocumentId = null;
            Map<String, Object> matchedDoc = null;
            
            if (documents != null) {
                for (Map<String, Object> doc : documents) {
                    String docName = (String) doc.get("name");
                    String ragDocId = (String) doc.get("id");
                    System.out.println("RAGFlow文档: '" + docName + "' (ID: " + ragDocId + ")");
                    
                    if (attach.getDocName().equals(docName)) {
                        ragflowDocumentId = ragDocId;
                        matchedDoc = doc;
                        System.out.println("✓ 找到匹配的文档，ID: " + ragflowDocumentId);
                        break;
                    }
                }
            }
            
            if (ragflowDocumentId == null) {
                System.err.println("✗ 未找到匹配的文档");
                System.err.println("请检查RAGFlow中是否有该文档");
                return R.fail("未找到对应的RAGFlow文档，本地文档名: " + attach.getDocName());
            }
            
            System.out.println("=====================================");
            
            // RAGFlow的chunks API需要特殊的认证环境，暂时返回文档信息和chunk数量
            // 用户可以在RAGFlow Web界面中查看完整的chunks
            Map<String, Object> result = new HashMap<>();
            result.put("documentId", ragflowDocumentId);
            result.put("datasetName", knowledgeInfoList.get(0).getKname());
            result.put("message", "知识片段统计数据");
            
            // 从文档列表中获取chunk_count（已经在获取文档列表时包含了）
            if (matchedDoc != null) {
                Integer chunkCount = (Integer) matchedDoc.get("chunk_count");
                result.put("chunkCount", chunkCount != null ? chunkCount : 0);
                result.put("documentInfo", matchedDoc);
                result.put("ragflowUrl", "http://192.168.0.150:7777/chunk/parsed/chunks?id=" + ragflowDocumentId + "&doc_id=" + ragflowDocumentId);
            }
            
            System.out.println("准备返回给前端的数据: " + result);
            System.out.println("返回的chunkCount: " + result.get("chunkCount"));
            System.out.println("返回的documentInfo: " + result.get("documentInfo"));
            
            return R.ok(result);
        } catch (Exception e) {
            return R.fail("获取RAGFlow片段失败: " + e.getMessage());
        }
    }

    /**
     * 上传文件到Ragflow
     */
    @PostMapping("/ragflow/upload/{kid}")
    public R<String> uploadToRagflow(@PathVariable String kid, @RequestParam("file") MultipartFile file) {
        boolean success = knowledgeInfoService.uploadToRagflow(kid, file);
        if (success) {
            return R.ok("文件上传到Ragflow成功！");
        } else {
            return R.fail("文件上传到Ragflow失败！");
        }
    }

    /**
     * 检查Ragflow服务状态
     */
    @GetMapping("/ragflow/status")
    public R<Boolean> checkRagflowStatus() {
        boolean available = knowledgeInfoService.checkRagflowStatus();
        return R.ok(available);
    }
    
    /**
     * 获取所有Ragflow数据集
     */
    @GetMapping("/ragflow/datasets")
    public R<List<Map<String, Object>>> getAllRagflowDatasets() {
        return R.ok(knowledgeInfoService.getAllRagflowDatasets());
    }
    
    /**
     * 从RAGFlow同步数据集到本地知识库
     */
    @PostMapping("/ragflow/sync")
    @Log(title = "同步RAGFlow数据集", businessType = BusinessType.INSERT)
    public R<String> syncFromRagflow() {
        try {
            int count = knowledgeInfoService.syncRagflowDatasetsToLocal();
            return R.ok("成功同步 " + count + " 个知识库", "" + count);
        } catch (Exception e) {
            return R.fail("同步失败: " + e.getMessage());
        }
    }

}
